# Minimal ONNX test models for Zhulong constant folding & DCE
# Only: const_chain_f32.onnx and dead_code_f32.onnx
# This script was created by ChatGPT 5.0.

import onnx
from onnx import helper, TensorProto


def save_model(graph, name, opset=13):
    model = helper.make_model(
        graph,
        opset_imports=[helper.make_operatorsetid("", opset)],
    )
    onnx.checker.check_model(model)
    onnx.save(model, name)
    print(f"=> wrote {name}")


# Constant chain (should fully fold to single constant output)
# z = (2.0 + 3.0) * 4.0 = 20.0
def make_const_chain_f32():
    Z = helper.make_tensor_value_info("z", TensorProto.FLOAT, [1])

    c2 = helper.make_tensor("c2", TensorProto.FLOAT, [1], [2.0])
    c3 = helper.make_tensor("c3", TensorProto.FLOAT, [1], [3.0])
    c4 = helper.make_tensor("c4", TensorProto.FLOAT, [1], [4.0])

    add0 = helper.make_node("Add", ["c2", "c3"], ["t_add"], name="add_const")
    mul0 = helper.make_node("Mul", ["t_add", "c4"], ["z"], name="mul_const")

    g = helper.make_graph([add0, mul0], "const_chain_f32", [], [Z],
                          initializer=[c2, c3, c4])
    return g, "const_chain_f32.onnx"


# Dead code elimination test
# live path: x + y -> z
# dead path: u + v (no consumer)
def make_dead_code_f32():
    X = helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 2])
    Y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 2])
    U = helper.make_tensor_value_info("u", TensorProto.FLOAT, [2, 2])
    V = helper.make_tensor_value_info("v", TensorProto.FLOAT, [2, 2])
    Z = helper.make_tensor_value_info("z", TensorProto.FLOAT, [2, 2])

    add_live = helper.make_node("Add", ["x", "y"], ["z"], name="add_live")
    add_dead = helper.make_node("Add", ["u", "v"], ["t_dead"], name="add_dead")

    g = helper.make_graph([add_live, add_dead], "dead_code_f32",
                          [X, Y, U, V], [Z])
    return g, "dead_code_f32.onnx"


if __name__ == "__main__":
    makers = [make_const_chain_f32, make_dead_code_f32]
    for mk in makers:
        g, fname = mk()
        save_model(g, fname, opset=13)
